import json
import shutil
import re
import os

from zhipuai import ZhipuAI

api_key = "4956db3da95d43eab77efa6937882679.Xm27NhBWZPJE3ORb"
client = ZhipuAI(api_key=api_key)

# 第一阶段LtM_CD阶段提示词及输出结果
get_latest_email_CD_input = "请帮我查询mailbox邮箱里最新一封邮件的内容."
get_latest_email_pi = "当前需求可以作为函数的参数的是:1.查看谁的邮箱"

get_latest_email_message_CD = [{"role":"user","content":get_latest_email_CD_input},
                               {"role":"assistant","content":get_latest_email_pi}]

# 第一阶段LtM_CM阶段提示词及输出结果
get_latest_email_CM_input = get_latest_email_CD_input + get_latest_email_pi
get_latest_email_description = "帮我写一个python函数，用于查看我的mailbox邮箱中的最新一封邮件信息，函数要求如下:\
                                1.函数参数userid，userid是字符串参数，默认只是'me'，表示查看我的邮件\
                                2.函数的返回结果是一个包含最新一封邮件信息的对象，返回的结果本身是一个json对象\
                                3.请将全部功能封装在一个函数内\
                                4.请在函数的编写过程中，在函数内部加入中文编写的函数详细说明文档，用于说明函数的功能，函数参数的情况以及函数返回结果等"

get_latest_email_message_CM = [{"role":"user","content":get_latest_email_CM_input},
                                {"role":"assistant","content":get_latest_email_description}]

# 第二阶段提示词及输出结果
# 将函数写入本地
with open(r"F:\ai\03大模型开发实战\08 低代码智能开发平台案例实践\tested functions\%s\%s.py" % ('get_latest_email','get_latest_email'), "r", encoding="utf8") as f:
    get_latest_email_function = f.read()

get_latest_email_message_FN = [{"role":"user","content":get_latest_email_description},
                               {"role":"assistant","content":get_latest_email_function}]

get_latest_email_prompt = {"stage1_CD":get_latest_email_message_CD,
                           "stage1_CM":get_latest_email_message_CM,
                           "stage2":get_latest_email_message_FN}

with open(r"F:\ai\03大模型开发实战\08 低代码智能开发平台案例实践\tested functions\%s\%s_prompt.json" % ('get_latest_email','get_latest_email'), "w", encoding="utf8") as f:
    json.dump(get_latest_email_prompt,f)

# with open(r"F:\ai\03大模型开发实战\08 低代码智能开发平台案例实践\%s.py" % 'get_latest_email_prompt', "r", encoding="utf8") as f:
#     get_latest_email_prompt_mes = f.read()
# print(get_latest_email_prompt_mes)

get_email_count_CD = "请帮我查看邮箱的未读的邮件数量."
get_email_count_pi = "当前需求可以作为函数的参数的是:\
                        1.查看谁的邮箱\
                        2.邮件的状态"
get_email_count_message_CD = [{"role":"user","content":get_email_count_CD},
                              {"role":"assistant","content":get_email_count_pi}]

get_email_count_CM_input = get_email_count_CD+get_email_count_pi
get_email_count_description = "帮我编写一个python函数，用于查看我的邮箱里有多少封未读邮件，函数的功能要求如下:\
                                1.函数有2个参数，分别是userid和status\
                                2.函数的参数userid表示查看谁的邮件，userid是一个字符串类型的参数，默认值是'me'，表示查看我的邮件\
                                3.函数的参数status表示查看未读的邮件，status是一个字符串类型的参数，默认值是'未读',表示查看未读的邮件\
                                4.函数的返回结果是一个包含邮件是否未读的对象列表数据，返回结果本身是一个json格式的对象\
                                5.请将全部功能封装在一个函数内\
                                6.请在函数的编写过程中，在函数内部加入中文编写的详细说明文档，用于说明函数的功能，函数的参数情况以及函数的返回结果等"
get_email_count_message_CM = [{"role":"user","content":get_email_count_CM_input},
                              {"role":"assistant","content":get_email_count_description}]

with open(r"F:\ai\03大模型开发实战\08 低代码智能开发平台案例实践\tested functions\%s\%s.py" % ('get_email_count','get_email_count'), "r", encoding="utf8") as f:
    get_email_count_function = f.read()

get_email_count_message_FN = [{"role":"user","content":get_email_count_description},
                              {"role":"assistant","content":get_email_count_function}]

get_email_count_prompt = {"stage1_CD":get_email_count_message_CD,
                          "stage1_CM":get_email_count_message_CM,
                          "stage2":get_email_count_message_FN}

with open(r"F:\ai\03大模型开发实战\08 低代码智能开发平台案例实践\tested functions\%s\%s_prompt.json" % ('get_email_count','get_email_count'), "w", encoding="utf8") as f:
    json.dump(get_email_count_prompt,f)

# with open(r"F:\ai\03大模型开发实战\08 低代码智能开发平台案例实践\%s.py" % 'get_latest_email_prompt', "r", encoding="utf8") as f:
#     get_latest_email_prompt_mes = f.read()
# print(get_latest_email_prompt_mes)

retrieve_emails_CD_input = "请帮我查看下我的邮箱中的前2封邮件"
retrieve_emails_pi = "当前需求中可以用作函数参数的是:\
                        1.查看谁的邮件\
                        2.查看多少封邮件"
retrieve_emails_message_CD = [{"role":"user","content":retrieve_emails_CD_input},
                              {"role":"assistant","content":retrieve_emails_pi}]

retrieve_emails_CM_input = retrieve_emails_CD_input+retrieve_emails_pi
retrieve_emails_description = "请帮我编写一个python函数，用于查看我的邮箱中的前几封邮件，该函数的功能要求如下：\
                                1.函数包含2个参数，分别是userid和n\
                                2.函数参数userid表示查看谁的邮件，userid是一个字符串类型的参数，默认值是'me'，表示查看我的邮件\
                                3.函数参数n表示查看几封邮件，n是一个整数类型的参数，默认值是1，表示查看几封邮件\
                                4.请将全部功能封装在一个函数内\
                                5.在函数编写过程中，在函数内部加入中文编写的详细文档说明，用于说明函数的功能、函数的参数情况和函数的返回结果"
retrieve_emails_message_CM = [{"role":"user","content":retrieve_emails_CM_input},
                              {"role":"assistant","content":retrieve_emails_description}]


with open(r"F:\ai\03大模型开发实战\08 低代码智能开发平台案例实践\tested functions\%s\%s.py" % ('retrieve_emails','retrieve_emails'), "r", encoding="utf8") as f:
    retrieve_emails_function = f.read()

retrieve_emails_message_FN = [{"role":"user","content":retrieve_emails_description},
                              {"role":"assistant","content":retrieve_emails_function}]

retrieve_emails_prompt = {"stage1_CD":retrieve_emails_message_CD,
                          "stage1_CM":retrieve_emails_message_CM,
                          "stage2":retrieve_emails_message_FN}
with open(r"F:\ai\03大模型开发实战\08 低代码智能开发平台案例实践\tested functions\%s\%s_prompt.py" % ('retrieve_emails','retrieve_emails'), "w", encoding="utf8") as f:
    json.dump(retrieve_emails_prompt,f)

with open(r"F:\ai\03大模型开发实战\08 低代码智能开发平台案例实践\资料\my_mail数据字典.md","r",encoding="utf8") as f:
    md_content = f.read()

system_messages = {"system_message_CD":[{"role":"system","content":md_content}],
                   "system_message_CM":[{"role":"system","content":md_content}],
                   "system_message":[{"role":"system","content":md_content}]}

with open(r"F:\ai\03大模型开发实战\08 低代码智能开发平台案例实践\tested functions\%s.json" % 'system_messages', "w", encoding="utf8") as f:
    json.dump(system_messages,f)

def remove_to_tested(function_name):
    """
    将函数同名文件夹由untested文件夹转移至tested文件夹内。\
    完成转移则说明函数通过测试，可以使用。
    """
    with open(r"F:\ai\03大模型开发实战\08 低代码智能开发平台案例实践\untested functions\%s\%s.py" % (function_name,function_name ), "r",encoding="utf8") as f:
        function_code = f.read()

    # 源文件夹路径
    src_dir = r"F:\ai\03大模型开发实战\08 低代码智能开发平台案例实践\untested functions"

    # 目标文件夹路径
    dst_dir = r"F:\ai\03大模型开发实战\08 低代码智能开发平台案例实践\tested functions"

    shutil.move(src_dir,dst_dir)

def extract_function_code(s,detail=0,tested=False,g=globals()):
    """
    函数提取函数，同时执行函数内容，可以选择打印函数信息，并选择代码保存的地址
    """

    def extract_code(s):
        """
        如果输入的字符串s是一个包含Python代码的Markdown格式字符串，提取出代码部分。
        否则，返回原字符串。

        参数:
        s: 输入的字符串。

        返回:
        提取出的代码部分，或原字符串。
        """
        # 判断字符串是否是Markdown格式
        if "```python" in s or "Python" in s or "PYTHON" in s:
            # 找到代码块的开始和结束位置
            if "import" in s:
                start = s.find("import")
            else:
                start = s.find("def")
            end = s.find("```",start)
            # 提取代码部分
            code = s[start:end]
        else:
            code = s
        return code

    # 提取代码字符串
    code = extract_code(s)

    # 提取函数名称
    match = re.search(r"def (\w+)",code)
    function_name = match.group(1)

    # 在untested文件夹内创建函数同名文件夹
    directory_untested = r"F:\ai\03大模型开发实战\08 低代码智能开发平台案例实践\untested functions\%s" % function_name
    if not os.path.exists(directory_untested):
        os.makedirs(directory_untested)


    # 将函数写入本地
    if tested == False:
        with open(r'F:\ai\03大模型开发实战\08 低代码智能开发平台案例实践\untested functions\%s/%s.py' % (function_name, function_name), 'w',encoding='utf-8') as f:
            f.write(code)
    else:
        # 调用remove_to_test函数将函数文件夹转移至tested文件夹内
        remove_to_tested(function_name)
        with open(r'F:\ai\03大模型开发实战\08 低代码智能开发平台案例实践\tested functions\%s/%s.py' % (function_name, function_name), 'w',encoding='utf-8') as f:
            f.write(code)

    # 执行该函数
    try:
        exec(code,g)
    except Exception as e:
        print("An error occurred while executing the code:")
        print(e)

    if detail == 0:
        print("The function name is:%s" % function_name)
    else:
        if tested == False:
            with open(r'F:\ai\03大模型开发实战\08 低代码智能开发平台案例实践\untested functions\%s/%s.py' % (function_name, function_name), 'r',encoding='utf-8') as f:
                content = f.read()
        else:
            with open(r'F:\ai\03大模型开发实战\08 低代码智能开发平台案例实践\tested functions\%s/%s.py' % (function_name, function_name), 'r',encoding='utf-8') as f:
                content = f.read()

        print(content)
    return function_name

def show_functions(tested=False,if_print=False):
    """
    打印tested或untested文件夹内全部函数
    """
    current_directory = ""
    if tested == False:
        current_directory = r"F:\ai\03大模型开发实战\08 低代码智能开发平台案例实践\untested functions"
    else:
        current_directory = r"F:\ai\03大模型开发实战\08 低代码智能开发平台案例实践\tested functions"

    files_and_directories = os.listdir(current_directory)
    # 过滤结果，只保留.py文件和非__pycache__文件夹
    files_and_directories = [name for name in files_and_directories if (os.path.splitext(name)[1] == ".py" or os.path.isdir(os.path.join(current_directory,name))) and name != ".ipynb_checkpoints"]

    if if_print != False:
        for name in files_and_directories:
            print(name)

    return files_and_directories

# files_and_directories = show_functions(tested=True,if_print=True)
# print(files_and_directories)


def code_generate(req,few_shot="all",model="glm-4",g=globals(),detail=0):
    """
    Function calling外部函数自动创建函数，可以根据用户的需求，直接将其翻译为Chat模型可以直接调用的外部函数代码。
    :param req: 必要参数，字符串类型，表示输入的用户需求；
    :param few_shot: 可选参数，默认取值为字符串all，用于描述Few-shot提示示例的选取方案，当输入字符串all时，则代表提取当前外部函数库中全部测试过的函数作为Few-shot；\
    而如果输入的是一个包含了多个函数名称的list，则表示使用这些函数作为Few-shot。
    :param model: 可选参数，表示调用的Chat模型，默认选取glm-4；
    :param g: 可选参数，表示extract_function_code函数作用域，默认为globals()，即在当前操作空间全域内生效；
    :param detail: 可选参数，默认取值为0，还可以取值为1，表示extract_function_code函数打印新创建的外部函数细节；
    :return：新创建的函数名称。需要注意的是，在函数创建时，该函数也会在当前操作空间被定义，后续可以直接调用；
    """

    # 提取提示示例的函数名称
    if few_shot == "all":
        few_shot_functions_name = show_functions(tested=True)
    elif type(few_shot) == list:
        few_shot_functions_name = few_shot

    # 读取各阶段系统提示
    with open(r"F:\ai\03大模型开发实战\08 低代码智能开发平台案例实践\tested functions\%s.json" % 'system_messages', "r",encoding="utf8") as f:
        system_messages = json.load(f)

    # 各阶段提示message对象
    few_shot_messages_CM = []
    few_shot_messages_CD = []
    few_shot_messages = []

    # 先保存第一条消息，也就是system message
    few_shot_messages_CD += system_messages["system_message_CD"]
    few_shot_messages_CM += system_messages["system_message_CM"]
    few_shot_messages += system_messages["system_message"]

    # 创建不同阶段提示message
    for function_name in few_shot_functions_name:
        with open(r"F:\ai\03大模型开发实战\08 低代码智能开发平台案例实践\tested functions\%s\%s_prompt.json" % (function_name,function_name),"r", encoding="utf8") as f:
            msg = json.load(f)
        few_shot_messages_CD += msg["stage1_CD"]
        few_shot_messages_CM += msg["stage1_CM"]
        few_shot_messages += msg["stage2"]

    # 读取用户需求，作为第一阶段CD环节User content
    new_req_CD_input = req
    few_shot_messages_CD.append({"role":"user","content":new_req_CD_input})
    print(few_shot_messages_CD)

    print('第一阶段CD环节提示创建完毕，正在进行CD提示...')

    # 第一阶段CD环节Chat模型调用过程
    resp_CD = client.chat.completions.create(model=model,messages=few_shot_messages_CD)
    new_req_pi = resp_CD.choices[0].message.content
    print(new_req_pi)

    print('第一阶段CD环节提示完毕')

    # 第一阶段CM环节Messages创建
    new_req_CM_input = new_req_CD_input+new_req_pi
    few_shot_messages_CM.append({"role":"user","content":new_req_CM_input})

    print('第一阶段CM环节提示创建完毕，正在进行第一阶段CM提示...')
    resp_CM = client.chat.completions.create(model=model,messages=few_shot_messages_CM)
    new_req_description = resp_CM.choices[0].message.content
    print(new_req_description)

    print('第一阶段CM环节提示完毕')

    # 第二阶段Messages创建过程
    few_shot_messages.append({"role":"user","content":new_req_description})

    print('第二阶段提示创建完毕，正在进行第二阶段提示...')
    resp_function = client.chat.completions.create(model=model,messages=few_shot_messages)

    new_req_function = resp_function.choices[0].message.content
    print(new_req_function)

    print('第二阶段提示完毕，准备运行函数并编写提示示例')

    # 提取函数并运行，创建函数名称对象，统一都写入untested文件夹内
    new_req_function_name = extract_function_code(s=new_req_function,detail=detail,g=g)
    print(r'新函数保存在F:\ai\03大模型开发实战\08 低代码智能开发平台案例实践\untested functions\%s\%s.py文件中' % (new_req_function_name, new_req_function_name))

    # 创建该函数提示示例
    new_req_message_CD = [{"role":"user","content":new_req_CD_input},
                          {"role":"assistant","content":new_req_pi}]

    new_req_message_CM = [{"role":"user","content":new_req_CM_input},
                          {"role":"assistant","content":new_req_description}]

    with open(r"F:\ai\03大模型开发实战\08 低代码智能开发平台案例实践\untested functions\%s\%s.py" % (new_req_function_name, new_req_function_name), "r", encoding="utf8") as f:
        new_req_function_code = f.read()

    new_req_message = [{"role":"user","content":new_req_description},
                       {"role":"assistant","content":new_req_function_code}]

    new_req_prompt = {"stage1_CD":new_req_message_CD,
                      "stage1_CM":new_req_message_CM,
                      "stage2":new_req_message}

    with open(r"F:\ai\03大模型开发实战\08 低代码智能开发平台案例实践\untested functions\%s\%s_prompt.json" % (new_req_function_name, new_req_function_name), 'w') as f:
        json.dump(new_req_prompt, f)

    print(r'新函数提示示例保存在F:\ai\03大模型开发实战\08 低代码智能开发平台案例实践\untested functions\%s\%s_prompt.json文件中' % (new_req_function_name, new_req_function_name))
    print('done')
    return new_req_function_name

req = "请查下我的邮箱里是否有来自陆小凤的未读邮件，有的话请解读下这封未读邮件的内容。"
few_shot_functions = ['get_latest_email', 'get_email_count']
function_name = code_generate(req=req, few_shot=few_shot_functions)
print(function_name)